/**
 * @file select_v2_custom_tiling.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include <cassert>
#include <cstddef>
#include <algorithm>
#include "dot_custom_tiling.h"

// bfloat16, float16, float, int8, int16, int32
constexpr uint32_t DATA_TYPE_SIZE[] = {2, 2, 4, 1, 2, 4};
constexpr uint32_t BLOCK_SIZE       = 32;
constexpr uint32_t BUFFER_NUM       = 2;
constexpr uint32_t UB_BLOCK_NUM     = 100;        // UB 最大可用 block 数量
constexpr uint32_t MAX_AVAILABLE_UB_BLOCK_NUM = UB_BLOCK_NUM / BUFFER_NUM * BUFFER_NUM;

// 计算单核内部 tile 参数
void TilingParamsCalc(uint32_t length, uint32_t alignNum,
                      uint32_t &tileNum, uint32_t &tileLength, uint32_t &lastTileLength)
{
    assert(alignNum != 0U);
    tileNum = length / (alignNum * MAX_AVAILABLE_UB_BLOCK_NUM);

    if ((length / alignNum) % MAX_AVAILABLE_UB_BLOCK_NUM == 0U || tileNum == 0U)
    {
        if (tileNum == 0U)
            tileNum = 1U;

        if (length < MAX_AVAILABLE_UB_BLOCK_NUM * alignNum)
        {
            tileLength = ((length / alignNum + BUFFER_NUM - 1) / BUFFER_NUM) * BUFFER_NUM * alignNum;
            lastTileLength = tileLength;
        }
        else
        {
            tileLength = MAX_AVAILABLE_UB_BLOCK_NUM * alignNum;
            lastTileLength = length - (tileNum - 1) * tileLength;
        }
    }
    else
    {
        tileNum++;
        tileLength = MAX_AVAILABLE_UB_BLOCK_NUM * alignNum;
        lastTileLength = length - (tileNum - 1) * tileLength;
    }
}

// 主 tiling 生成函数
void GenerateTilingData(uint8_t *tilingBuf, uint32_t blockDim)
{
    DotCustomTilingData *tiling = reinterpret_cast<DotCustomTilingData *>(tilingBuf);

    const uint32_t totalLength = tiling->totalLength;
    const uint32_t dataTypeSize = DATA_TYPE_SIZE[tiling->dataType];
    const uint32_t alignNum = BLOCK_SIZE / dataTypeSize;

    assert(alignNum != 0U && blockDim != 0U);

    // 对齐后的总长度
    uint32_t totalLengthAligned = (totalLength % alignNum == 0U)
                                      ? totalLength
                                      : ((totalLength + alignNum - 1) / alignNum) * alignNum;

    /* ---------- 1. 可均分场景 ---------- */
    if ((totalLengthAligned / alignNum) % blockDim == 0U)
    {
        const uint32_t blockLength = totalLengthAligned / blockDim;

        uint32_t tileNum, tileLength, lastTileLength;
        TilingParamsCalc(blockLength, alignNum, tileNum, tileLength, lastTileLength);

        tiling->blockLength   = blockLength;
        tiling->tileNum       = tileNum;
        tiling->tileLength    = tileLength;
        tiling->lastTileLength= lastTileLength;
        tiling->isEvenCore    = 1U;
    }
    /* ---------- 2. 不可均分场景 ---------- */
    else
    {
        // 每个核“平均”处理的元素数
        const uint32_t avgElem = totalLengthAligned / blockDim;

        // 需要多承担 1 个 alignNum 的核数
        const uint32_t formerNum = totalLengthAligned % blockDim;
        const uint32_t tailNum   = blockDim - formerNum;

        // 向上 / 向下对齐后实际分到的长度
        const uint32_t formerLength = ((avgElem + alignNum - 1) / alignNum) * alignNum;
        const uint32_t tailLength   = (avgElem / alignNum) * alignNum;

        // 检查总和是否一致
        assert(formerLength * formerNum + tailLength * tailNum == totalLengthAligned);

        uint32_t formerTileNum, formerTileLen, formerLastTileLen;
        uint32_t tailTileNum,   tailTileLen,   tailLastTileLen;

        TilingParamsCalc(formerLength, alignNum,
                         formerTileNum, formerTileLen, formerLastTileLen);
        TilingParamsCalc(tailLength,   alignNum,
                         tailTileNum,   tailTileLen,   tailLastTileLen);

        tiling->isEvenCore              = 0U;
        tiling->formerNum               = formerNum;
        tiling->formerLength            = formerLength;
        tiling->formerTileNum           = formerTileNum;
        tiling->formerTileLength        = formerTileLen;
        tiling->formerLastTileLength    = formerLastTileLen;

        tiling->tailNum                 = tailNum;
        tiling->tailLength              = tailLength;
        tiling->tailTileNum             = tailTileNum;
        tiling->tailTileLength          = tailTileLen;
        tiling->tailLastTileLength      = tailLastTileLen;
    }
}